package com.atguigu.bigdata.spark.sql

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, LongType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

/**
  * create by undeRdoG on  2021-06-19  12:37
  * 凡心所向，素履以往，生如逆旅，一苇以航。
  */
object Spark03_SparkSQL_UDAF {

  def main(args: Array[String]): Unit = {
    val sparkConf = new SparkConf().setMaster("local[*]").setAppName("SparkSQL")

    val spark = SparkSession.builder().config(sparkConf).getOrCreate()
    import spark.implicits._


    val df: DataFrame = spark.read.json("datas/user.json")

    df.createOrReplaceTempView("user")

    spark.udf.register("AgeAvg",new MyAvgUDAF())

    spark.sql("select AgeAvg(age) from user").show()
  }


  /**
    * UDAF
    * 自定义聚合函数类：计算年龄平均值
    * 1  继承UserDefinedAggregateFunction
    **/

  class MyAvgUDAF extends UserDefinedAggregateFunction {
    // 输入数据结构
    override def inputSchema: StructType = {
      StructType(
        Array(
          StructField("age", LongType)
        )
      )
    }

    //  缓冲区数据结构
    override def bufferSchema: StructType = {
      StructType(
        Array(
          StructField("total", LongType),
          StructField("count", LongType)
        )
      )
    }

    //  计算结果的输出类型
    override def dataType: DataType = {
      LongType
    }

    //  函数的稳定性
    override def deterministic: Boolean = true

    // 缓冲区初始化
    override def initialize(buffer: MutableAggregationBuffer): Unit = {
      buffer(0) = 0L
      buffer(1) = 0L
    }

    // 根据输入值 来更新缓冲区
    override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
      buffer.update(0,buffer.getLong(0) + input.getLong(0))
      buffer.update(1,buffer.getLong(1) + 1)

    }


    //  缓冲区数据合并
    override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
      buffer1.update(0,buffer2.getLong(0) + buffer1.getLong(0))

      buffer1.update(1,buffer2.getLong(1) + buffer1.getLong(1))
    }

    //  计算
    override def evaluate(buffer: Row): Any = {

      buffer.getLong(0) / buffer.getLong(1)

    }
  }

}
